"""
公共频道WebSocket订阅
包含所有公共数据频道的订阅功能
"""
import asyncio
from typing import Callable, Dict, List, Optional
from loguru import logger

try:
    from okx.websocket.WsPublicAsync import WsPublicAsync
    from okx.websocket.WsPrivateAsync import WsPrivateAsync
except ImportError:
    logger.warning("okx SDK未安装，请运行: pip install okx")
    WsPublicAsync = None
    WsPrivateAsync = None


class PublicChannelsManager:
    """公共频道订阅管理器（包含产品频道和12个公共频道）"""
    
    # 产品类型
    INST_TYPES = {
        "SPOT": "币币",
        "MARGIN": "币币杠杆",
        "SWAP": "永续合约",
        "FUTURES": "交割合约",
        "OPTION": "期权",
    }
    
    # WebSocket URL
    WS_PUBLIC_URL = "wss://ws.okx.com:8443/ws/v5/public"
    WS_PUBLIC_SIMULATE_URL = "wss://wspap.okx.com:8443/ws/v5/public"
    WS_BUSINESS_URL = "wss://ws.okx.com:8443/ws/v5/business"
    WS_BUSINESS_SIMULATE_URL = "wss://wspap.okx.com:8443/ws/v5/business"
    
    def __init__(self, sandbox: bool = True):
        """
        初始化公共频道订阅管理器
        
        Args:
            sandbox: 是否使用模拟盘
        """
        if WsPublicAsync is None:
            raise ImportError("okx SDK未安装，请运行: pip install okx")
        
        self.sandbox = sandbox
        self.public_url = self.WS_PUBLIC_SIMULATE_URL if sandbox else self.WS_PUBLIC_URL
        self.business_url = self.WS_BUSINESS_SIMULATE_URL if sandbox else self.WS_BUSINESS_URL
        
        # WebSocket连接
        self.ws_public: Optional[WsPublicAsync] = None
        self.ws_business: Optional[WsPublicAsync] = None
        self.is_public_connected = False
        self.is_business_connected = False
        
        # 订阅管理
        self.subscriptions: Dict[str, Dict] = {}
        
        logger.info(f"公共频道订阅管理器已初始化 (模式: {'模拟盘' if sandbox else '实盘'})")
    
    async def connect_public(self):
        """连接到公共WebSocket服务器"""
        if self.is_public_connected:
            return
        
        self.ws_public = WsPublicAsync(url=self.public_url)
        await self.ws_public.start()
        self.is_public_connected = True
        logger.info("公共WebSocket连接成功")
    
    async def connect_business(self):
        """连接到业务WebSocket服务器"""
        if self.is_business_connected:
            return
        
        self.ws_business = WsPublicAsync(url=self.business_url)
        await self.ws_business.start()
        self.is_business_connected = True
        logger.info("业务WebSocket连接成功")
    
    async def disconnect(self):
        """断开所有WebSocket连接"""
        for key in list(self.subscriptions.keys()):
            await self._unsubscribe_by_key(key)
        
        self.is_public_connected = False
        self.is_business_connected = False
        logger.info("WebSocket连接已断开")
    
    def _default_callback(self, message: Dict):
        """默认回调函数"""
        logger.info(f"收到推送: {message}")
    
    # 0. 产品频道（instruments channel）
    async def subscribe_instruments(self, inst_type: str, inst_family: Optional[str] = None, callback: Optional[Callable] = None):
        """
        订阅产品频道 - 产品状态变化时推送
        
        Args:
            inst_type: 产品类型 (SPOT/MARGIN/SWAP/FUTURES/OPTION)
            inst_family: 交易品种（可选），如 BTC-USD
            callback: 回调函数
        """
        await self.connect_public()
        
        # 验证产品类型
        if inst_type not in self.INST_TYPES:
            raise ValueError(f"无效的产品类型: {inst_type}，支持的类型: {list(self.INST_TYPES.keys())}")
        
        args = {"channel": "instruments", "instType": inst_type}
        if inst_family:
            args["instFamily"] = inst_family
        
        callback = callback or self._default_callback
        await self.ws_public.subscribe([args], callback=callback)
        
        key = f"instruments_{inst_type}_{inst_family}" if inst_family else f"instruments_{inst_type}"
        self.subscriptions[key] = {"args": [args], "callback": callback, "ws": "public"}
        
        inst_type_name = self.INST_TYPES.get(inst_type, inst_type)
        logger.info(f"已订阅产品频道: {inst_type_name}")
        if inst_family:
            logger.info(f"  交易品种: {inst_family}")
    
    async def subscribe_all_spot(self, callback: Optional[Callable] = None):
        """订阅所有现货产品"""
        await self.subscribe_instruments("SPOT", callback=callback)
    
    async def subscribe_all_swap(self, callback: Optional[Callable] = None):
        """订阅所有永续合约"""
        await self.subscribe_instruments("SWAP", callback=callback)
    
    async def subscribe_all_futures(self, callback: Optional[Callable] = None):
        """订阅所有交割合约"""
        await self.subscribe_instruments("FUTURES", callback=callback)
    
    async def subscribe_btc_instruments(self, callback: Optional[Callable] = None):
        """订阅BTC相关产品"""
        await self.subscribe_instruments("SWAP", "BTC-USD", callback)
        await self.subscribe_instruments("FUTURES", "BTC-USD", callback)
    
    # 1.1 持仓总量频道
    async def subscribe_open_interest(self, inst_id: str, callback: Optional[Callable] = None):
        """订阅持仓总量频道 - 每3s推送一次"""
        await self.connect_public()
        args = [{"channel": "open-interest", "instId": inst_id}]
        callback = callback or self._default_callback
        await self.ws_public.subscribe(args, callback=callback)
        self.subscriptions[f"open-interest_{inst_id}"] = {"args": args, "callback": callback, "ws": "public"}
        logger.info(f"已订阅持仓总量频道: {inst_id}")
    
    # 1.2 资金费率频道
    async def subscribe_funding_rate(self, inst_id: str, callback: Optional[Callable] = None):
        """订阅资金费率频道 - 30-90秒推送一次"""
        await self.connect_public()
        args = [{"channel": "funding-rate", "instId": inst_id}]
        callback = callback or self._default_callback
        await self.ws_public.subscribe(args, callback=callback)
        self.subscriptions[f"funding-rate_{inst_id}"] = {"args": args, "callback": callback, "ws": "public"}
        logger.info(f"已订阅资金费率频道: {inst_id}")
    
    # 1.3 限价频道
    async def subscribe_price_limit(self, inst_id: str, callback: Optional[Callable] = None):
        """订阅限价频道 - 有变化时每200ms推送"""
        await self.connect_public()
        args = [{"channel": "price-limit", "instId": inst_id}]
        callback = callback or self._default_callback
        await self.ws_public.subscribe(args, callback=callback)
        self.subscriptions[f"price-limit_{inst_id}"] = {"args": args, "callback": callback, "ws": "public"}
        logger.info(f"已订阅限价频道: {inst_id}")
    
    # 1.4 期权定价频道
    async def subscribe_option_summary(self, inst_family: str, callback: Optional[Callable] = None):
        """订阅期权定价频道 - 一次性推送所有"""
        await self.connect_public()
        args = [{"channel": "opt-summary", "instFamily": inst_family}]
        callback = callback or self._default_callback
        await self.ws_public.subscribe(args, callback=callback)
        self.subscriptions[f"opt-summary_{inst_family}"] = {"args": args, "callback": callback, "ws": "public"}
        logger.info(f"已订阅期权定价频道: {inst_family}")
    
    # 1.5 预估交割/行权/结算价格频道
    async def subscribe_estimated_price(self, inst_type: str, inst_family: str, callback: Optional[Callable] = None):
        """订阅预估价格频道 - 交割前一小时开始推送"""
        await self.connect_public()
        args = [{"channel": "estimated-price", "instType": inst_type, "instFamily": inst_family}]
        callback = callback or self._default_callback
        await self.ws_public.subscribe(args, callback=callback)
        self.subscriptions[f"estimated-price_{inst_type}_{inst_family}"] = {"args": args, "callback": callback, "ws": "public"}
        logger.info(f"已订阅预估价格频道: {inst_type} {inst_family}")
    
    # 1.6 标记价格频道
    async def subscribe_mark_price(self, inst_id: str, callback: Optional[Callable] = None):
        """订阅标记价格频道 - 有变化时每200ms推送，无变化每10s推送"""
        await self.connect_public()
        args = [{"channel": "mark-price", "instId": inst_id}]
        callback = callback or self._default_callback
        await self.ws_public.subscribe(args, callback=callback)
        self.subscriptions[f"mark-price_{inst_id}"] = {"args": args, "callback": callback, "ws": "public"}
        logger.info(f"已订阅标记价格频道: {inst_id}")
    
    # 1.7 指数行情频道
    async def subscribe_index_tickers(self, inst_id: str, callback: Optional[Callable] = None):
        """订阅指数行情频道 - 有变化时每100ms推送，否则1分钟推送"""
        await self.connect_public()
        args = [{"channel": "index-tickers", "instId": inst_id}]
        callback = callback or self._default_callback
        await self.ws_public.subscribe(args, callback=callback)
        self.subscriptions[f"index-tickers_{inst_id}"] = {"args": args, "callback": callback, "ws": "public"}
        logger.info(f"已订阅指数行情频道: {inst_id}")
    
    # 1.8 标记价格K线频道
    async def subscribe_mark_price_candle(self, inst_id: str, channel: str = "mark-price-candle1D", callback: Optional[Callable] = None):
        """订阅标记价格K线频道 - 最快每秒推送"""
        await self.connect_business()
        args = [{"channel": channel, "instId": inst_id}]
        callback = callback or self._default_callback
        await self.ws_business.subscribe(args, callback=callback)
        self.subscriptions[f"{channel}_{inst_id}"] = {"args": args, "callback": callback, "ws": "business"}
        logger.info(f"已订阅标记价格K线频道: {channel} {inst_id}")
    
    # 1.9 指数K线频道
    async def subscribe_index_candle(self, inst_id: str, channel: str = "index-candle30m", callback: Optional[Callable] = None):
        """订阅指数K线频道 - 最快每秒推送"""
        await self.connect_business()
        args = [{"channel": channel, "instId": inst_id}]
        callback = callback or self._default_callback
        await self.ws_business.subscribe(args, callback=callback)
        self.subscriptions[f"{channel}_{inst_id}"] = {"args": args, "callback": callback, "ws": "business"}
        logger.info(f"已订阅指数K线频道: {channel} {inst_id}")
    
    # 1.10 平台公共爆仓单频道
    async def subscribe_liquidation_orders(self, inst_type: str, callback: Optional[Callable] = None):
        """订阅爆仓单频道"""
        await self.connect_public()
        args = [{"channel": "liquidation-orders", "instType": inst_type}]
        callback = callback or self._default_callback
        await self.ws_public.subscribe(args, callback=callback)
        self.subscriptions[f"liquidation-orders_{inst_type}"] = {"args": args, "callback": callback, "ws": "public"}
        logger.info(f"已订阅爆仓单频道: {inst_type}")
    
    # 1.11 自动减仓预警频道
    async def subscribe_adl_warning(self, inst_type: str, inst_family: str, callback: Optional[Callable] = None):
        """订阅自动减仓预警频道 - 普通状态每分钟推送，预警状态每秒推送"""
        await self.connect_public()
        args = [{"channel": "adl-warning", "instType": inst_type, "instFamily": inst_family}]
        callback = callback or self._default_callback
        await self.ws_public.subscribe(args, callback=callback)
        self.subscriptions[f"adl-warning_{inst_type}_{inst_family}"] = {"args": args, "callback": callback, "ws": "public"}
        logger.info(f"已订阅自动减仓预警频道: {inst_type} {inst_family}")
    
    async def _unsubscribe_by_key(self, key: str):
        """取消订阅"""
        if key not in self.subscriptions:
            return
        
        sub_info = self.subscriptions[key]
        ws_type = sub_info["ws"]
        ws = self.ws_public if ws_type == "public" else self.ws_business
        
        await ws.unsubscribe(sub_info["args"], callback=sub_info["callback"])
        del self.subscriptions[key]
        logger.info(f"已取消订阅: {key}")
    
    def get_subscriptions(self) -> List[str]:
        """获取当前所有订阅"""
        return list(self.subscriptions.keys())


class EconomicCalendarChannel:
    """经济日历频道（需要API认证，仅实盘，VIP1+）"""
    
    WS_BUSINESS_URL = "wss://ws.okx.com:8443/ws/v5/business"
    
    def __init__(self, api_key: str, secret_key: str, passphrase: str):
        """
        初始化经济日历频道
        
        Args:
            api_key: API密钥
            secret_key: 密钥
            passphrase: 口令
        """
        if WsPrivateAsync is None:
            raise ImportError("okx SDK未安装，请运行: pip install okx")
        
        self.api_key = api_key
        self.secret_key = secret_key
        self.passphrase = passphrase
        self.ws: Optional[WsPrivateAsync] = None
        self.is_connected = False
        
        logger.info("经济日历频道已初始化（需API认证，仅实盘，VIP1+）")
    
    async def connect(self):
        """连接到WebSocket服务器"""
        if self.is_connected:
            return
        
        self.ws = WsPrivateAsync(
            apiKey=self.api_key,
            passphrase=self.passphrase,
            secretKey=self.secret_key,
            url=self.WS_BUSINESS_URL,
            useServerTime=False
        )
        await self.ws.start()
        self.is_connected = True
        logger.info("经济日历WebSocket连接成功")
    
    async def subscribe(self, callback: Optional[Callable] = None):
        """订阅经济日历频道"""
        await self.connect()
        args = [{"channel": "economic-calendar"}]
        callback = callback or self._default_callback
        await self.ws.subscribe(args, callback=callback)
        logger.info("已订阅经济日历频道")
    
    async def unsubscribe(self, callback: Optional[Callable] = None):
        """取消订阅"""
        if not self.is_connected:
            return
        args = [{"channel": "economic-calendar"}]
        callback = callback or self._default_callback
        await self.ws.unsubscribe(args, callback=callback)
        logger.info("已取消订阅经济日历频道")
    
    def _default_callback(self, message: Dict):
        """默认回调函数"""
        logger.info(f"收到经济日历推送: {message}")
